import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import pandas as pd
import numpy as np

def plot_predictions(test_dates: pd.Series, y_test: np.ndarray, 
                    predictions: np.ndarray, title: str = '用电量预测结果') -> None:
    """绘图预测结果
    
    Args:
        test_dates: 测试集日期
        y_test: 实际值
        predictions: 预测值
        title: 标题
    """
    plt.figure(figsize=(12, 6))
    
    # 按日期排序
    sort_idx = np.argsort(test_dates)
    test_dates = test_dates.iloc[sort_idx]
    y_test = y_test[sort_idx]
    predictions = predictions[sort_idx]
    
    plt.scatter(test_dates, y_test, color='blue', label='实际值', alpha=0.6)
    plt.plot(test_dates, y_test, color='blue', linestyle='-', alpha=0.3)
    plt.scatter(test_dates, predictions, color='red', label='预测值', alpha=0.6)
    plt.plot(test_dates, predictions, color='red', linestyle='--', alpha=0.3)
    
    # 设置图表格式
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    plt.gcf().autofmt_xdate()
    plt.rcParams['font.sans-serif'] = ['SimHei']    # Windows
    
    plt.xlabel('年月')
    plt.ylabel('用电量')
    plt.title(title)
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    
def plot_model_comparison(model_scores: dict) -> None:
    """绘图-模型R2对比
    
    Args:
        model_scores: 模型名称和得分的字典
    """
    plt.figure(figsize=(10, 6))
    names = list(model_scores.keys())
    scores = list(model_scores.values())
    
    plt.bar(names, scores)
    plt.title('模型性能对比')
    plt.xlabel('模型')
    plt.ylabel('R2得分')
    plt.xticks(rotation=45)
    plt.tight_layout() 